from __future__ import print_function
import torch
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F
from .submodule import *
import math
import timm

# 模块父类
class SubModule(nn.Module):
    def __init__(self):
        super(SubModule, self).__init__()

    # 权重初始化
    def weight_init(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.Conv3d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

# 特征提取，提取多尺度特征
class Feature(SubModule):
    def __init__(self):
        super(Feature, self).__init__()
        pretrained =  True
        model = timm.create_model('mobilenetv2_100', pretrained=pretrained, features_only=True) # 从timm下载mobilenetv2_100结构和权重
        layers = [1,2,3,5,6]
        chans = [16, 24, 32, 96, 160]

        # mobilenetv2的特征提取部分(前3层)
        self.conv_stem = model.conv_stem # 卷积层conv_stem
        self.bn1 = model.bn1 # 批归一化层bn1
        self.act1 = model.act1 # ReLU激活函数层act1

        # mobilenetv2的blocks块部分(第0~5层)
        self.block0 = torch.nn.Sequential(*model.blocks[0:layers[0]]) # 解包blocks块，取块内的第0层
        self.block1 = torch.nn.Sequential(*model.blocks[layers[0]:layers[1]]) # 取第1层
        self.block2 = torch.nn.Sequential(*model.blocks[layers[1]:layers[2]]) # 取第2层
        self.block3 = torch.nn.Sequential(*model.blocks[layers[2]:layers[3]]) # 取第3、4层
        self.block4 = torch.nn.Sequential(*model.blocks[layers[3]:layers[4]]) # 取第5层

        self.deconv32_16 = Conv2x(chans[4], chans[3], deconv=True, concat=True) # 拼接块，没有调用

    def forward(self, x):
        x = self.act1(self.bn1(self.conv_stem(x))) # 输入3通道，输出32通道，卷积核3*3，步长1，1/2原尺寸
        x2 = self.block0(x) # 输入32通道，输出16通道，尺寸不变
        x4 = self.block1(x2) # 输入16通道，输出24通道，1/4原尺寸
        # return x4,x4,x4,x4
        x8 = self.block2(x4) # 输入24通道，输出32通道，1/8原尺寸
        x16 = self.block3(x8) # 输入32通道，输出96通道，1/16原尺寸
        x32 = self.block4(x16) # 输入96通道，输出160通道，1/32原尺寸
        return [x4, x8, x16, x32] # 返回多尺度特征，1/4、1/8、1/16、1/32原尺寸，24、32、96、160通道

# 特征融合上采样，小尺寸向大尺寸融合(1/4原尺寸特征是最终特征，同时输出所有中间产物的不同尺寸特征)
class FeatUp(SubModule):
    def __init__(self):
        super(FeatUp, self).__init__()
        chans = [16, 24, 32, 96, 160]
        self.deconv32_16 = Conv2x(chans[4], chans[3], deconv=True, concat=True) # 拼接块，对x：输入160通道，输出96通道，2d转置尺寸*2；在通道维度拼接张量x,rem
        self.deconv16_8 = Conv2x(chans[3]*2, chans[2], deconv=True, concat=True) # 拼接块，对x：输入192通道，输出32通道，2d转置尺寸*2；在通道维度拼接张量x,rem
        self.deconv8_4 = Conv2x(chans[2]*2, chans[1], deconv=True, concat=True) # 拼接块，对x：输入64通道，输出24通道，2d转置尺寸*2；在通道维度拼接张量x,rem
        self.conv4 = BasicConv(chans[1]*2, chans[1]*2, kernel_size=3, stride=1, padding=1) # 基本卷积块，2d卷积，输入48通道，输出48通道，尺寸不变

        self.weight_init() # 权重初始化

    def forward(self, featL, featR=None):
        x4, x8, x16, x32 = featL # 多尺度特征，1/4、1/8、1/16、1/32原尺寸
        y4, y8, y16, y32 = featR

        x16 = self.deconv32_16(x32, x16) # 拼接块，拼接左视图1/32和1/16原尺寸特征为1/16原尺寸特征，输出192通道
        y16 = self.deconv32_16(y32, y16)
        x8 = self.deconv16_8(x16, x8) # 拼接块，拼接左视图1/16和1/8原尺寸特征为1/8原尺寸特征，输出64通道
        y8 = self.deconv16_8(y16, y8)
        x4 = self.deconv8_4(x8, x4) # 拼接块，拼接左视图1/8和1/4原尺寸特征为1/4原尺寸特征，输出48通道
        y4 = self.deconv8_4(y8, y4)
        x4 = self.conv4(x4) # 基本卷积块，通道不变，尺寸不变，输出48通道
        y4 = self.conv4(y4)

        return [x4, x8, x16, x32], [y4, y8, y16, y32] # 返回左右视图多尺度特征，1/4、1/8、1/16、1/32原尺寸，48、64、192、160通道

# 上下文与几何信息融合CGF模块
class Context_Geometry_Fusion(SubModule):
    def __init__(self, cv_chan, im_chan): # 成本体积通道，特征通道，输出张量形状与输入的成本体积一致
        super(Context_Geometry_Fusion, self).__init__()

        self.semantic = nn.Sequential(
            BasicConv(im_chan, im_chan//2, kernel_size=1, stride=1, padding=0), # 基本卷积块，2d卷积，通道/2，尺寸不变
            nn.Conv2d(im_chan//2, cv_chan, 1) # 2d卷积，改变通道，尺寸不变
            )

        self.att = nn.Sequential(
            BasicConv(cv_chan, cv_chan, is_3d=True, bn=True, relu=True, kernel_size=(1,5,5), padding=(0,2,2), stride=1, dilation=1), # 基本卷积块，3d卷积，通道不变，尺寸不变
            nn.Conv3d(cv_chan, cv_chan, kernel_size=1, stride=1, padding=0, bias=False) # 3d卷积，通道不变，尺寸不变
            )

        self.agg = BasicConv(cv_chan, cv_chan, is_3d=True, bn=True, relu=True, kernel_size=(1,5,5), padding=(0,2,2), stride=1, dilation=1) # 基本卷积块，3d卷积，通道不变，尺寸不变

        self.weight_init() # 权重初始化

    def forward(self, cv, feat): # 输入成本体积(几何信息)，特征(上下文信息)
        '''
        '''
        feat = self.semantic(feat).unsqueeze(2) # 处理特征，输出形状与成本体积一致
        att = self.att(feat+cv) # 空间注意力张量(编码了要注意的信息)，成本体积和特征相加，5*5卷积
        cv = torch.sigmoid(att)*feat + cv # 加性注意力，空间注意力权重*特征(上下文信息)，融合成本体积(几何信息)，*表示对应元素相乘(哈达玛积)
        cv = self.agg(cv) # 5*5卷积

        return cv # 返回张量形状与输入的成本体积一致

# 成本聚合，沙漏融合结构
class hourglass_fusion(nn.Module):
    def __init__(self, in_channels): # 输入8通道
        super(hourglass_fusion, self).__init__()

        self.conv1 = nn.Sequential(
            BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 基本卷积块，3d卷积，输入8通道，输出16通道，尺寸/2
            BasicConv(in_channels*2, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1) # 基本卷积块，3d卷积，通道不变，尺寸不变
            )

        self.conv2 = nn.Sequential(
            BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 基本卷积块，3d卷积，输入16通道，输出32通道，尺寸/2
            BasicConv(in_channels*4, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1) # 基本卷积块，3d卷积，通道不变，尺寸不变
            )

        self.conv3 = nn.Sequential(
            BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 基本卷积块，3d卷积，输入32通道，输出48通道，尺寸/2
            BasicConv(in_channels*6, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1) # 基本卷积块，3d卷积，通道不变，尺寸不变
            )

        self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True, relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 基本卷积块，3d转置卷积，输入48通道，输出32通道，尺寸*2

        self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True, relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 基本卷积块，3d转置卷积，输入32通道，输出16通道，尺寸*2

        self.conv1_up = BasicConv(in_channels*2, 1, deconv=True, is_3d=True, bn=False, relu=False, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 基本卷积块，3d转置卷积，无批归一化和LeakyReLU，输入16通道，输出1通道，尺寸*2

        self.agg_0 = nn.Sequential(
            BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1), # 基本卷积块，3d卷积，输入64通道，输出32通道，尺寸不变
            BasicConv(in_channels*4, in_channels*4, is_3d=True, kernel_size=3, padding=1, stride=1), # 基本卷积块，3d卷积，通道不变，尺寸不变
            BasicConv(in_channels*4, in_channels*4, is_3d=True, kernel_size=3, padding=1, stride=1) # 同上
            )

        self.agg_1 = nn.Sequential(
            BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1), # 基本卷积块，3d卷积，输入32通道，输出16通道，尺寸不变
            BasicConv(in_channels*2, in_channels*2, is_3d=True, kernel_size=3, padding=1, stride=1), # 基本卷积块，3d卷积，通道不变，尺寸不变
            BasicConv(in_channels*2, in_channels*2, is_3d=True, kernel_size=3, padding=1, stride=1) # 同上
            )

        self.CGF_32 = Context_Geometry_Fusion(in_channels*6, 160) # 上下文与几何信息融合CGF模块，输入成本体积通道，特征通道，输出张量形状与输入的成本体积一致
        self.CGF_16 = Context_Geometry_Fusion(in_channels*4, 192) # 上下文与几何信息融合CGF模块
        self.CGF_8 = Context_Geometry_Fusion(in_channels*2, 64) # 上下文与几何信息融合CGF模块

    def forward(self, x, imgs): # 输入注意力特征体积(B,8,D/4,H/4,W/4)，上采样融合多尺度特征(96、64、192、160通道)
        conv1 = self.conv1(x) # 3d卷积下采样，输入8通道，输出16通道，尺寸/2，(B,16,D/8,H/8,W/8)
        conv2 = self.conv2(conv1) # 3d卷积下采样，输入16通道，输出32通道，尺寸/2，(B,32,D/16,H/16,W/16)
        conv3 = self.conv3(conv2) # 3d卷积下采样，输入32通道，输出48通道，尺寸/2，(B,48,D/32,H/32,W/32)

        conv3 = self.CGF_32(conv3, imgs[3]) # 上下文与几何信息融合CGF模块，输入成本体积(几何信息)(B,48,D/32,H/32,W/32)，1/32原尺寸左特征(上下文信息)(B,160,H/32,W/32)，输出(B,48,D/32,H/32,W/32)
        conv3_up = self.conv3_up(conv3) # 3d转置卷积上采样，输入48通道，输出32通道，尺寸*2，(B,32,D/16,H/16,W/16)

        conv2 = torch.cat((conv3_up, conv2), dim=1) # 在通道维度拼接上采样张量和下采样张量，(B,64,D/16,H/16,W/16)
        conv2 = self.agg_0(conv2) # 3d卷积减少通道到拼接前，输入64通道，输出32通道，(B,32,D/16,H/16,W/16)

        conv2 = self.CGF_16(conv2, imgs[2]) # 上下文与几何信息融合CGF模块，输入成本体积(几何信息)(B,32,D/16,H/16,W/16)，1/16原尺寸左特征(上下文信息)(B,192,H/16,W/16)，输出(B,32,D/16,H/16,W/16)
        conv2_up = self.conv2_up(conv2) # 3d转置卷积上采样，输入32通道，输出16通道，尺寸*2，(B,16,D/8,H/8,W/8)

        conv1 = torch.cat((conv2_up, conv1), dim=1) # 在通道维度拼接上采样张量和下采样张量，(B,32,D/8,H/8,W/8)
        conv1 = self.agg_1(conv1) # 3d卷积减少通道到拼接前，输入32通道，输出16通道，(B,16,D/8,H/8,W/8)

        conv1 = self.CGF_8(conv1, imgs[1]) # 上下文与几何信息融合CGF模块，输入成本体积(几何信息)(B,16,D/8,H/8,W/8)，1/8原尺寸左特征(上下文信息)(B,64,H/8,W/8)，输出(B,16,D/8,H/8,W/8)
        conv = self.conv1_up(conv1) # 3d转置卷积上采样，输入16通道，输出1通道，尺寸*2，(B,1,D/4,H/4,W/4)

        return conv

class CGI_Stereo(nn.Module):
    def __init__(self, maxdisp):
        super(CGI_Stereo, self).__init__()
        self.maxdisp = maxdisp
        self.feature = Feature() # 特征提取，使用mobilenetv2的特征提取器
        self.feature_up = FeatUp() # 特征融合上采样，小尺寸向大尺寸融合(1/4原尺寸特征是最终特征，同时输出所有中间产物的不同尺寸特征)
        chans = [16, 24, 32, 96, 160]

        self.stem_2 = nn.Sequential(
            BasicConv(3, 32, kernel_size=3, stride=2, padding=1), # 基本卷积块，2d卷积，批归一化，LeakyReLU，输入3通道，输出32通道，尺寸/2
            nn.Conv2d(32, 32, 3, 1, 1, bias=False), # 2d卷积，通道不变，尺寸不变
            nn.BatchNorm2d(32), nn.ReLU()
            )
        self.stem_4 = nn.Sequential(
            BasicConv(32, 48, kernel_size=3, stride=2, padding=1), # 基本卷积块，2d卷积，批归一化，LeakyReLU，输入32通道，输出48通道，尺寸/2
            nn.Conv2d(48, 48, 3, 1, 1, bias=False), # 2d卷积，通道不变，尺寸不变
            nn.BatchNorm2d(48), nn.ReLU()
            )

        self.spx = nn.Sequential(nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1)) # 2d转置卷积，输入64通道，输出9通道，尺寸*2
        self.spx_2 = Conv2x(32, 32, True) # 拼接块，对x：输入32通道，输出32通道，2d转置尺寸*2；在通道维度拼接张量x,rem
        self.spx_4 = nn.Sequential(
            BasicConv(96, 32, kernel_size=3, stride=1, padding=1), # 基本卷积块，2d卷积，批归一化，LeakyReLU，输入96通道，输出32通道，尺寸不变
            nn.Conv2d(32, 32, 3, 1, 1, bias=False), # 2d卷积，通道不变，尺寸不变
            nn.BatchNorm2d(32), nn.ReLU()
            )

        self.conv = BasicConv(96, 48, kernel_size=3, padding=1, stride=1) # 基本卷积块，2d卷积，批归一化，LeakyReLU，输入96通道，输出48通道，尺寸不变
        self.desc = nn.Conv2d(48, 48, kernel_size=1, padding=0, stride=1) # 2d卷积，通道不变，尺寸不变
        self.corr_stem = BasicConv(1, 8, is_3d=True, kernel_size=3, stride=1, padding=1) # 基本卷积块，3d卷积，批归一化，LeakyReLU，输入1通道，输出8通道，尺寸不变
        self.semantic = nn.Sequential(
            BasicConv(96, 32, kernel_size=3, stride=1, padding=1), # 基本卷积块，2d卷积，批归一化，LeakyReLU，输入96通道，输出32通道，尺寸不变
            nn.Conv2d(32, 8, kernel_size=1, padding=0, stride=1, bias=False) # 2d卷积，输入32通道，输出8通道，尺寸不变
            )
        self.agg = BasicConv(8, 8, is_3d=True, kernel_size=(1,5,5), padding=(0,2,2), stride=1) # 基本卷积块，3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变
        self.hourglass_fusion = hourglass_fusion(8) # 成本聚合，沙漏融合结构，输入8通道

    def forward(self, left, right):
        features_left = self.feature(left) # 多尺度特征(ImageNet上预训练的mobilenetv2_100特征提取器)，返回多尺度特征，1/4、1/8、1/16、1/32原尺寸，24、32、96、160通道
        features_right = self.feature(right)
        features_left, features_right = self.feature_up(features_left, features_right) # 上采样融合多尺度特征，返回左右视图多尺度特征，1/4、1/8、1/16、1/32原尺寸，48、64、192、160通道

        stem_2x = self.stem_2(left) # 1/2原尺寸特征，对原图像两次2d卷积+批归一化+LeakyReLU(ReLU)，32通道，(B,32,H/2,W/2)
        stem_4x = self.stem_4(stem_2x) # 1/4原尺寸特征，对1/2原尺寸特征两次2d卷积+批归一化+LeakyReLU(ReLU)，48通道，(B,48,H/4,W/4)
        stem_2y = self.stem_2(right)
        stem_4y = self.stem_4(stem_2y)

        features_left[0] = torch.cat((features_left[0], stem_4x), 1) # 1/4原尺寸上下文特征，在通道维度连接1/4原尺寸的上采样融合特征和提取的特征，96通道，(B,96,H/4,W/4)
        features_right[0] = torch.cat((features_right[0], stem_4y), 1)

        match_left = self.desc(self.conv(features_left[0])) # 用于构建相关性体积的1/4原尺寸上下文特征，两次2d卷积+批归一化+LeakyReLU，48通道，(B,48,H/4,W/4)
        match_right = self.desc(self.conv(features_right[0]))

        corr_volume = build_norm_correlation_volume(match_left, match_right, self.maxdisp//4) # 构建归一化相关性体积，输入1/4原尺寸左右上下文特征(48通道)，最大视差/4，输出(B,1,D/4,H/4,W/4)
        corr_volume = self.corr_stem(corr_volume) # 相关性体积，3d卷积+批归一化+LeakyReLU，输出8通道，尺寸不变，(B,8,D/4,H/4,W/4)
        feat_volume = self.semantic(features_left[0]).unsqueeze(2) # 左特征体积，1/4原尺寸左特征，两次2d卷积+批归一化+LeakyReLU，增加第2维度(视差)，输出8通道，尺寸不变，(B,8,1,H/4,W/4)
        volume = self.agg(feat_volume * corr_volume) # 注意力特征体积，左特征体积(广播机制扩充视差维度)*相关性体积，3d卷积，输出(B,8,D/4,H/4,W/4)

        cost = self.hourglass_fusion(volume, features_left) # 成本聚合，沙漏融合结构，输入注意力特征体积，上采样融合多尺度特征(96、64、192、160通道)，(B,1,D/4,H/4,W/4)

        xspx = self.spx_4(features_left[0]) # 对1/4原尺寸上下文特征两次2d卷积+批归一化+LeakyReLU(ReLU)，输出32通道，尺寸不变，(B,32,H/4,W/4)
        xspx = self.spx_2(xspx, stem_2x) # 拼接块，对xspx：通道不变，2d转置尺寸*2；在通道维度拼接张量xspx,stem_2x，输出64通道，1/2原尺寸，(B,64,H/2,W/2)
        spx_pred = self.spx(xspx) # 2d转置卷积，输入64通道，输出9通道，尺寸*2(原尺寸)，(B,9,H,W)
        spx_pred = F.softmax(spx_pred, 1) # 原尺寸"超像素"权重，在通道维度做softmax，输出(B,9,H,W)

        disp_samples = torch.arange(0, self.maxdisp//4, dtype=cost.dtype, device=cost.device) # 创建0~D/4的视差序列
        disp_samples = disp_samples.view(1, self.maxdisp//4, 1, 1).repeat(cost.shape[0], 1, cost.shape[3], cost.shape[4]) # 视差序列体积，改变形状为(1,D/4,1,1)，重复视差序列为(B,D/4,H/4,W/4)
        pred = regression_topk(cost.squeeze(1), disp_samples, 2) # 1/4原尺寸回归视差，在视差维度取两个相关性最大值，输入聚合的成本体积(B,D/4,H/4,W/4)，视差序列体积(B,D/4,H/4,W/4)，2，输出(B,1,H/4,W/4)
        pred_up = context_upsample(pred, spx_pred) # 上下文上采样，输入1/4原尺寸回归视差(1通道)，原尺寸"超像素"权重(9通道)，输出(B,H,W)(1/4尺寸视差)

        if self.training: # 训练模式
            return [pred_up*4, pred.squeeze(1)*4] # 返回全尺寸视差List，[上下文上采样视差(B,H,W),1/4原尺寸回归视差(B,H/4,W/4)]
        else: # 验证模式
            return [pred_up*4] # 返回全尺寸视差List，[上下文上采样视差(B,H,W)]
